import cv2
import numpy as np
import time
from typing import Union, Generator

# 自定义识别颜色


def set_color():
    # 蓝色
    lower_blue = np.array([100, 60, 100], dtype=np.uint8)
    upper_blue = np.array([124, 230, 230], dtype=np.uint8)

    # 绿色
    lower_green = np.array([35, 60, 100], dtype=np.uint8)
    upper_green = np.array([77, 230, 230], dtype=np.uint8)
    import os

    print(os.getcwd())
    # 读取自定义的颜色
    with open('./data/color_custom_good.txt', 'r') as color_good:
        colors = color_good.readlines()
        if (colors):
            color1 = colors[0].split(",")

            color1[0] = int(color1[0])-5
            color1[1] = 60
            color1[2] = 100
            lower_good = np.array(color1, dtype=np.uint8)

            color1[0] += 10
            color1[1] = 230
            color1[2] = 230
            upper_good = np.array(color1, dtype=np.uint8)
        else:
            # 默认绿色为good
            lower_good = lower_green
            upper_good = upper_green
    with open('./data/color_custom_bad.txt', 'r') as color_bad:
        colors = color_bad.readlines()
        if (colors):
            color2 = colors[0].split(",")

            color2[0] = int(color2[0])-5
            color2[1] = 60
            color2[2] = 100

            lower_bad = np.array(color2, dtype=np.uint8)

            color2[0] += 10
            color2[1] = 230
            color2[2] = 230
            upper_bad = np.array(color2, dtype=np.uint8)
        else:
            # 蓝色为bad
            lower_bad = lower_blue
            upper_bad = upper_blue
    #print(lower_good, upper_good, lower_bad, upper_bad)
    return lower_good, upper_good, lower_bad, upper_bad
    # return lower_green, upper_green, lower_blue, upper_blue


def getHue(r, g, b):
    r, g, b = r / 255, g / 255, b / 255
    mx = max(r, g, b)
    mn = min(r, g, b)
    m = mx - mn
    if mx == mn:
        h = 0
    elif mx == r:
        if g > b:
            h = (g-b)*60/m
        else:
            h = (g-b)*60/m + 360
    elif mx == g:
        h = (b-r)*60/m + 120
    elif mx == b:
        h = (r-g)*60/m + 240
    return h / 2

def findend(cap,framenum,lower,upper,frameRate):
    num=framenum-frameRate//2
    while(num<framenum):
        cap.set(cv2.CAP_PROP_POS_FRAMES, num)

        ret, frame = cap.read()

        if ret:
            frame = cv2.flip(frame, 1)
            kernel = np.ones((3, 3), np.uint8)
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            mask = cv2.inRange(hsv, lower, upper)

            # extrapolate the hand to fill dark spots within
            mask = cv2.dilate(mask, kernel, iterations=4)

            # blur the image
            mask = cv2.GaussianBlur(mask, (5, 5), 100)
            # 找轮廓
            cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]

            if len(cnts) > 0:
                c = max(cnts, key=cv2.contourArea)  # 找到图像中最大的轮廓
                ((x, y), radius) = cv2.minEnclosingCircle(c)  # 找到最小的圆包含这个轮廓，返回坐标和半径
                print(radius,num)
                if radius <= 40:
                    return num
            else:
                return num

        num += frameRate//8  # 4的时候效果较好，但速度较慢，可以考虑先检测到然后再具体检测
    return framenum

def findstart(cap,framenum,lower,upper,frameRate):
    num=framenum-frameRate//2
    while(num<framenum):
        cap.set(cv2.CAP_PROP_POS_FRAMES, num)

        ret, frame = cap.read()

        if ret:
            frame = cv2.flip(frame, 1)
            kernel = np.ones((3, 3), np.uint8)
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            mask = cv2.inRange(hsv, lower, upper)

            # extrapolate the hand to fill dark spots within
            mask = cv2.dilate(mask, kernel, iterations=4)

            # blur the image
            mask = cv2.GaussianBlur(mask, (5, 5), 100)
            # 找轮廓
            cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]

            if len(cnts) > 0:
                c = max(cnts, key=cv2.contourArea)  # 找到图像中最大的轮廓
                ((x, y), radius) = cv2.minEnclosingCircle(c)  # 找到最小的圆包含这个轮廓，返回坐标和半径
                print(radius,num)
                if radius > 40:
                    return num

        num += frameRate//8  # 4的时候效果较好，但速度较慢，可以考虑先检测到然后再具体检测
    return framenum

# 检测并剪裁视频
def detect_color(filepath: str, goodRGB: tuple[int, int, int], badRGB: tuple[int, int, int]) -> Generator[Union[int, list[dict[str, Union[str, float]]]], None, None]:
    data: list[dict[str, Union[str, float]]] = []
    cap = cv2.VideoCapture(filepath)
    #cap = cv2.VideoCapture(0)
    bad = 0  # bad手势区域
    good = 0  # good手势区域
    bc = 1  # 计数
    gc = 1  # 计数
    curMark: dict[str, Union[str, float]] = {}  # 当前标记

    time_start = time.time()

    totalFrame = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    frameRate = cap.get(cv2.CAP_PROP_FPS)

    # pauseTime = int(1000/frameRate+0.5) #四舍五入
    # print(pauseTime)

    framenum = -1

    # 获取设定的good和bad颜色
    # lower_good, upper_good, lower_bad, upper_bad = set_color()
    lower_good, upper_good, lower_bad, upper_bad = (
        np.array([getHue(*goodRGB) - 10, 60, 100]),
        np.array([getHue(*goodRGB) + 10, 230, 230]),
        np.array([getHue(*badRGB) - 10, 60, 100]),
        np.array([getHue(*badRGB) + 10, 230, 230]),
    )

    while (1):
        #print(framenum)
        ret, frame = cap.read()

        if ret:
            if (time_start == 0):
                time_start = time.time()
            frame = cv2.flip(frame, 1)
            kernel = np.ones((3, 3), np.uint8)

            # define region of interest
            #roi = frame[100:300, 100:300]

            #cv2.rectangle(frame, (100, 100), (300, 300), (0, 255, 0), 0)
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

            # # define range of skin color in HSV
            # lower_skin = np.array([0, 20, 70], dtype=np.uint8)
            # upper_skin = np.array([20, 255, 255], dtype=np.uint8)

            # extract skin colur imagw
            mask = cv2.inRange(hsv, lower_good, upper_good)

            # extrapolate the hand to fill dark spots within
            mask = cv2.dilate(mask, kernel, iterations=4)

            # blur the image
            mask = cv2.GaussianBlur(mask, (5, 5), 100)

            # 找轮廓
            cnts = cv2.findContours(
                mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
            #center = None

            # extract skin colur imagw
            mask2 = cv2.inRange(hsv, lower_bad, upper_bad)

            # extrapolate the hand to fill dark spots within
            mask2 = cv2.dilate(mask2, kernel, iterations=4)

            # blur the image
            mask2 = cv2.GaussianBlur(mask2, (5, 5), 100)

            # 找轮廓
            cnts2 = cv2.findContours(
                mask2.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
            #center2 = None

            radius = 0  # good部分的半径
            radius2 = 0  # bad部分的半径

            if len(cnts) > 0:
                c = max(cnts, key=cv2.contourArea)  # 找到图像中最大的轮廓
                ((x, y), radius) = cv2.minEnclosingCircle(c)  # 找到最小的圆包含这个轮廓，返回坐标和半径

                # M = cv2.moments(c)  # 计算图像的矩
                # center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))  # 通过矩计算图像的质心
            else:
                if good == 1:  # 如果屏幕中已经检测不到相应颜色，则结束该标记
                    good = 0
                    # mylog.write("\n <end>" + str(framenum / frameRate + 0.3) + "</end> \n </marks>")
                    # 以较小的帧率更细致地查找end帧
                    pos=findend(cap,framenum,lower_good,upper_good,frameRate)
                    print(pos,framenum)
                    curMark['end'] = pos / frameRate + 0.5
                    #curMark['end'] = framenum / frameRate + 0.5
                    data.append(curMark)
                    curMark = {}
            #cv2.namedWindow('frame', cv2.WINDOW_AUTOSIZE)

            if len(cnts2) > 0:

                c = max(cnts2, key=cv2.contourArea)  # 找到图像中最大的轮廓
                ((x, y), radius2) = cv2.minEnclosingCircle(c)  # 找到最小的圆包含这个轮廓，返回坐标和半径
                #
                # M = cv2.moments(c)  # 计算图像的矩
                # center2 = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))  # 通过矩计算图像的质心

            else:
                if bad == 1:  # #如果屏幕中已经检测不到相应颜色，则结束该标记
                    # mylog.write("\n <end>" + str(framenum / frameRate + 0.3) + "</end> \n </
                    # 以较小的帧率更细致地查找end帧
                    pos = findend(cap, framenum, lower_bad, upper_bad, frameRate)
                    print(pos, framenum)
                    curMark['end'] = pos / frameRate + 0.5
                    #curMark['end'] = framenum / frameRate + 0.5
                    data.append(curMark)
                    curMark = {}
                    bad = 0

            if radius > radius2:  # good标记的范围更明显，则识别为good
                # 判断以下他的半径，太小的话就认为是噪声，就忽略掉
                if radius > 40:
                    if bad == 0 and good == 0:
                        print("good", gc)
                        gc += 1
                        good = 1  # good手势开始

                        curMark['label'] = 'good'
                        pos=findstart(cap,framenum,lower_good,upper_good,frameRate)
                        curMark['start'] = pos / frameRate - 0.5
                        print('start',pos,framenum)

                    # cv2.circle(frame, (int(x), int(y)), int(radius), (0, 255, 255), 2)  # 在检测到的轮廓周围画圆圈
                    # cv2.circle(frame, center, 5, (0, 0, 255), -1)

                else:
                    if good == 1:
                        good = 0
                        pos = findend(cap, framenum, lower_good, upper_good, frameRate)
                        print(pos, framenum)
                        curMark['end'] = pos / frameRate + 0.5
                        #curMark['end'] = framenum / frameRate + 0.5
                        data.append(curMark)
                        curMark = {}

            else:  # bad标记的范围更明显，则识别为bad
                # 判断以下他的半径，太小的话就认为是噪声，就忽略掉
                if radius2 > 40:
                    if good == 0 and bad == 0:
                        bad = 1  # bad手势开始
                        print("bad", bc)
                        bc += 1
                        curMark['label'] = 'bad'
                        pos = findstart(cap, framenum, lower_bad, upper_bad, frameRate)
                        curMark['start'] = pos / frameRate - 0.5
                        print('start', pos, framenum)

                    # cv2.circle(frame, (int(x), int(y)), int(radius2), (0, 255, 255), 2)  # 在检测到的轮廓周围画圆圈
                    # cv2.circle(frame, center2, 5, (0, 0, 255), -1)

                else:
                    if bad == 1:  # bad手势结束
                        pos = findend(cap, framenum, lower_bad, upper_bad, frameRate)
                        print(pos, framenum)
                        curMark['end'] = pos / frameRate + 0.5
                        #curMark['end'] = framenum / frameRate + 0.5
                        data.append(curMark)
                        curMark = {}
                        bad = 0

            # 注释掉的部分是用于显示识别情况的
            # cv2.namedWindow('frame', cv2.WINDOW_AUTOSIZE)
            # cv2.imshow('mask', mask)
            # cv2.imshow('mask2', mask2)
            # cv2.imshow('frame', frame)
            # cv2.waitKey(5)
        # break
        else:
            break

        #原本在裁剪的时候要求保留最后的片段（即没有结束符的最后部分)，因此这里必须是-2而不是-1，至少保留1帧可供裁剪
        if(framenum<totalFrame-2 and framenum+frameRate>totalFrame):
             framenum = totalFrame-2
        else:
             framenum += frameRate // 1  # 4的时候效果较好，但速度较慢，可以考虑先检测到然后再具体检测
        #framenum+=1
        cap.set(cv2.CAP_PROP_POS_FRAMES, framenum)

        yield framenum*100//totalFrame
        # if(framenum%10==0):
        #     print(framenum*100//totalFrame)


    cv2.destroyAllWindows()
    cap.release()
    print(time.time() - time_start)
    yield data





